import os.path
import websocket  # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
from flask import Flask, request, jsonify
from gevent import pywsgi
from image import get_images
import random
import uuid
import json
from util import get_process_by_prompt

app = Flask(__name__)

# 本地ip 默认127.0.0.1:8188
server_address = "192.168.1.36:8188"
# 输出路径
current_path = os.getcwd()
savePath = os.path.join(current_path, 'output\\images')

print(savePath)

@app.route('/requestApi/createImage', methods=['POST'])
def createImage():
    client_id = str(uuid.uuid4())
    random_num = random.randint(1, 2147483647)
    file_path = "\\\\Nas-02\\\\常用软件\\\\研发一部\\\\研发一部\\\\model\\\\createImage.json"
    with open(file_path, 'r', encoding='utf-8') as file:
        prompt = json.loads(file.read())
        # 调节工作流参数
        prompt["25"]["inputs"]["noise_seed"] = random_num
        prompt["17"]["inputs"]["steps"] = 1
        prompt["28"]["inputs"]["string"] = "logo, ai, origin, turing"
        prompt["70"]["inputs"]["int"] = 128
        prompt["71"]["inputs"]["int"] = 128
        prompt["75"]["inputs"]["output_path"] = savePath
    # ws监听进度并输出图片信息
    ws = websocket.WebSocket()
    ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
    progress = 0
    for data in get_images(ws, prompt, client_id, server_address, savePath):
        if data['node'] is not None:
            p = get_process_by_prompt(prompt, data['node'])
            progress = p if p > progress else progress
        data["progress"] = progress
        print(data)
        yield json.dumps(data).encode('utf-8')
    ws.close()


if __name__ == "__main__":
    # 开启api
    server = pywsgi.WSGIServer(('0.0.0.0', 1588), app)
    server.serve_forever()
